from flask import Flask, render_template, request, jsonify, redirect, url_for
import torch
import torchvision.transforms as transforms
from PIL import Image
import os
import io
import base64
import sys
import json
import sqlite3
import datetime
import uuid
from flask import Flask, render_template, request, redirect, url_for, flash, jsonify, session

# 添加项目根目录到Python路径
# sys.path.append(r'c:\Users\Ipuuy\Desktop\project')
sys.path.append(r'C:\Users\Lenovo\PycharmProjects\DRnet\project')
from model import load_model

app = Flask(__name__)
app.secret_key = 'supersecretkey'  # 设置secret_key以使用flash消息

# 临时患者信息文件路径
TEMP_PATIENT_INFO_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'temp_patient_info.json')

# 初始化数据库
def init_db():
    conn = sqlite3.connect('detection_records.db')
    cursor = conn.cursor()
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS records (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        timestamp TEXT,
        filename TEXT,
        image_path TEXT,
        diagnosis TEXT,
        confidence REAL,
        all_probs TEXT,
        image_data BLOB
    )
    ''')
    
    # 添加测试数据，确保历史记录页面有内容显示
    # 检查是否已经存在测试记录
    cursor.execute("SELECT id FROM records WHERE image_path LIKE ?", ('uploads/%',))
    existing_records = cursor.fetchall()
    
    # 如果记录数量较少，添加一些有效的测试数据
    if len(existing_records) < 3:
        # 查找uploads目录中已有的文件
        import os
        upload_dir = os.path.join('static', 'uploads')
        if os.path.exists(upload_dir):
            existing_files = [f for f in os.listdir(upload_dir) if os.path.isfile(os.path.join(upload_dir, f))]
            
            # 如果有现有的文件，使用它们作为测试数据
            if existing_files:
                test_records = []
                conditions = ['轻度非增值性视网膜病变', '中度非增值性视网膜病变', '正常']
                
                for i, file_name in enumerate(existing_files[:3]):  # 最多使用3个文件
                    condition_index = i % len(conditions)
                    confidence = 0.7 + (i * 0.05)
                    
                    # 构造合适的概率分布
                    probs = [0.0] * 5
                    if condition_index == 0:  # 轻度
                        probs = [0.10, confidence, 0.07, 0.02, 0.01]
                    elif condition_index == 1:  # 中度
                        probs = [0.05, 0.20, confidence, 0.08, 0.02]
                    else:  # 正常
                        probs = [confidence, 0.10, 0.03, 0.01, 0.01]
                    
                    # 确保概率总和为1
                    total = sum(probs)
                    probs = [p / total for p in probs]
                    
                    image_path = f'uploads/{file_name}'
                    test_records.append((
                        datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                        file_name,
                        image_path,
                        conditions[condition_index],
                        confidence,
                        json.dumps(probs),
                        None
                    ))
                
                # 插入测试数据
                cursor.executemany("INSERT INTO records (timestamp, filename, image_path, diagnosis, confidence, all_probs, image_data) VALUES (?, ?, ?, ?, ?, ?, ?)", test_records)
    
    conn.commit()
    conn.close()

# 保存检测记录
def save_detection_record(filename, image_path, diagnosis, confidence, all_probs, image_data=None, patient_info=None):
    conn = sqlite3.connect('detection_records.db')
    cursor = conn.cursor()
    timestamp = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    all_probs_json = json.dumps(all_probs)
    
    # 尝试获取患者信息
    patient_name = ''
    patient_number = ''
    doctor_name = ''
    
    if patient_info:
        # 如果传入了患者信息，直接使用
        patient_name = patient_info.get('patient_name', '')
        patient_number = patient_info.get('patient_number', '')
        doctor_name = patient_info.get('doctor_name', '')
    else:
        # 尝试从临时文件中读取患者信息
        try:
            if os.path.exists(TEMP_PATIENT_INFO_FILE):
                with open(TEMP_PATIENT_INFO_FILE, 'r', encoding='utf-8') as f:
                    temp_info = json.load(f)
                    patient_name = temp_info.get('patient_name', '')
                    patient_number = temp_info.get('patient_number', '')
                    doctor_name = temp_info.get('doctor_name', '')
        except Exception as e:
            print(f"从文件读取患者信息失败: {e}")
    
    # 插入记录，包含患者信息
    cursor.execute(
        "INSERT INTO records (timestamp, filename, image_path, diagnosis, confidence, all_probs, image_data, patient_name, patient_number, doctor_name) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
        (timestamp, filename, image_path, diagnosis, confidence, all_probs_json, image_data, patient_name, patient_number, doctor_name)
    )
    
    record_id = cursor.lastrowid
    conn.commit()
    conn.close()
    return record_id

# 获取所有检测记录，支持搜索和分页
def get_all_records(search_term=None, page=1, per_page=10):
    # 使用绝对路径连接到正确的数据库文件
    db_path = r'c:\Users\Lenovo\PycharmProjects\DRnet\project\change\detection_records.db'
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    # 计算偏移量
    offset = (page - 1) * per_page
    
    # 确保表存在
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS records (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        timestamp TEXT,
        filename TEXT,
        image_path TEXT,
        diagnosis TEXT,
        confidence REAL,
        all_probs TEXT,
        image_data BLOB,
        patient_name TEXT,
        patient_number TEXT,
        doctor_name TEXT
    )
    ''')
    
    # 先获取总记录数
    if search_term:
        # 支持搜索患者姓名、患者编号、主治医生姓名和文件名
        cursor.execute("SELECT COUNT(*) FROM records WHERE patient_name LIKE ? OR patient_number LIKE ? OR doctor_name LIKE ? OR filename LIKE ?", 
                     (f"%{search_term}%", f"%{search_term}%", f"%{search_term}%", f"%{search_term}%"))
    else:
        cursor.execute("SELECT COUNT(*) FROM records")
    
    count_result = cursor.fetchone()
    total_count = count_result[0] if count_result else 0
    
    # 然后获取记录列表
    if search_term:
        # 支持搜索患者姓名、患者编号、主治医生姓名和文件名
        cursor.execute("SELECT id, timestamp, patient_name, diagnosis, doctor_name FROM records WHERE patient_name LIKE ? OR patient_number LIKE ? OR doctor_name LIKE ? OR filename LIKE ? ORDER BY timestamp DESC LIMIT ? OFFSET ?", 
                     (f"%{search_term}%", f"%{search_term}%", f"%{search_term}%", f"%{search_term}%", per_page, offset))
    else:
        # 普通查询并分页
        cursor.execute("SELECT id, timestamp, patient_name, diagnosis, doctor_name FROM records ORDER BY timestamp DESC LIMIT ? OFFSET ?", 
                     (per_page, offset))
    
    records = cursor.fetchall()
    
    conn.close()
    return records, total_count

# 删除检测记录
def delete_record(record_id):
    conn = sqlite3.connect('detection_records.db')
    cursor = conn.cursor()
    
    # 先获取记录信息，检查是否存在
    cursor.execute("SELECT filename, image_path FROM records WHERE id = ?", (record_id,))
    record = cursor.fetchone()
    
    if record:
        # 删除记录
        cursor.execute("DELETE FROM records WHERE id = ?", (record_id,))
        conn.commit()
        
        # 注意：这里不自动删除关联的图像文件，因为其他记录可能也在使用相同的文件
        # 如果确实需要删除文件，可以取消下面的注释
        # if record[1]:  # 如果有image_path
        #     image_path = os.path.join('static', record[1])
        #     if os.path.exists(image_path):
        #         os.remove(image_path)
    
    conn.close()
    return record is not None

# 获取单个检测记录详情
def get_record_detail(record_id):
    conn = sqlite3.connect('detection_records.db')
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM records WHERE id = ?", (record_id,))
    record = cursor.fetchone()
    conn.close()
    if record:
        try:
            # 转换all_probs从JSON字符串回到列表
            all_probs = json.loads(record[6])
        except json.JSONDecodeError:
            all_probs = [0.0, 0.0, 0.0, 0.0, 0.0]  # 默认值
        
        # 确保image_path格式正确
        image_path = record[3]
        filename = record[2]
        
        # 解决Windows路径分隔符问题 - 确保使用正斜杠
        if image_path:
            image_path = image_path.replace('\\', '/')
        
        # 智能处理图像路径
        if image_path:
            # 检查图像文件是否存在于static目录下
            full_path = os.path.join('static', image_path)
            
            # 如果文件不存在，尝试不同的路径组合
            if not os.path.exists(full_path):
                # 尝试各种可能的路径格式
                possible_paths = []
                
                # 情况1: 只有文件名，尝试放在uploads目录下
                if not '/' in image_path and not '\\' in image_path and filename:
                    possible_paths.append(f'uploads/{image_path}')
                
                # 情况2: 尝试直接使用文件名
                if filename:
                    possible_paths.append(filename)
                    possible_paths.append(f'uploads/{filename}')
                
                # 情况3: 尝试从路径中提取文件名
                if image_path:
                    basename = os.path.basename(image_path)
                    possible_paths.append(basename)
                    possible_paths.append(f'uploads/{basename}')
                
                # 检查所有可能的路径
                found_path = None
                for path in possible_paths:
                    test_path = os.path.join('static', path)
                    if os.path.exists(test_path):
                        found_path = path.replace('\\', '/')
                        break
                
                # 如果找到有效路径，更新image_path
                if found_path:
                    image_path = found_path
        
        record_dict = {
            'id': record[0],
            'timestamp': record[1],
            'filename': record[2],
            'image_path': image_path,
            'diagnosis': record[4],
            'confidence': record[5],
            'all_probs': all_probs,
            'image_data': record[7],
            'patient_name': record[8] if len(record) > 8 else '',
            'patient_number': record[9] if len(record) > 9 else '',
            'doctor_name': record[10] if len(record) > 10 else ''
        }
        return record_dict
    return None

# 在应用启动时初始化数据库
init_db()

# 确保中文显示正常
import matplotlib

matplotlib.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]

# 配置上传文件夹
UPLOAD_FOLDER = 'static/uploads'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# 确保上传文件夹存在
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# 尝试加载模型
try:
    model = load_model()
    print("模型加载成功")
except Exception as e:
    print(f"模型加载失败: {e}")
    # 使用模拟模型作为备用
    model = None
    print("使用模拟模型作为备用")


# 图像预处理函数
def preprocess_image(image):
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])
    ])
    return transform(image).unsqueeze(0)


# 预测函数
def predict(image):
    # 定义5种糖尿病视网膜病变等级
    disease_classes = [
        "正常",
        "轻度非增值性视网膜病变",
        "中度非增值性视网膜病变",
        "重度非增值性视网膜病变",
        "增值性视网膜病变"
    ]

    if model is None:
        # 模拟预测结果
        import random
        # 生成模拟概率分布（确保和为1）
        import numpy as np
        probs = np.random.dirichlet(np.ones(5), size=1)[0]
        pred = np.argmax(probs)
        return disease_classes[pred], float(probs[pred]), [float(p) for p in probs]

    input_tensor = preprocess_image(image)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    input_tensor = input_tensor.to(device)

    with torch.no_grad():
        output = model(input_tensor)
        prob = torch.softmax(output, dim=1)
        pred = torch.argmax(prob, dim=1).item()
        confidence = prob[0, pred].item()
        # 获取所有类别的概率
        all_probs = [float(p) for p in prob[0]]

        return disease_classes[pred], confidence, all_probs


# 检查文件类型是否允许
def allowed_file(filename):
    return '.' in filename and \
        filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


# 首页路由
@app.route('/')
def home():
    # 尝试从session或临时文件中获取患者信息
    current_patient = None
    
    # 1. 先尝试从session获取
    if 'patient_info' in session:
        current_patient = session['patient_info']
    # 2. 如果session中没有，尝试从临时文件获取
    elif os.path.exists(TEMP_PATIENT_INFO_FILE):
        try:
            with open(TEMP_PATIENT_INFO_FILE, 'r', encoding='utf-8') as f:
                current_patient = json.load(f)
        except Exception as e:
            print(f"读取临时患者信息失败: {e}")
    
    return render_template('index.html', current_patient=current_patient)

# 患者信息输入页面路由
@app.route('/patient_info')
def patient_info():
    return render_template('patient_info.html')

# 保存患者信息的路由
@app.route('/save_patient_info', methods=['POST'])
def save_patient_info():
    patient_name = request.form.get('patient_name')
    patient_number = request.form.get('patient_number')
    doctor_name = request.form.get('doctor_name')
    
    # 验证输入
    if not patient_name or not patient_number or not doctor_name:
        flash('请填写所有必填信息')
        return redirect(url_for('patient_info'))
    
    # 保存患者信息到session
    session['patient_info'] = {
        'patient_name': patient_name,
        'patient_number': patient_number,
        'doctor_name': doctor_name
    }
    
    # 同时保存到临时文件，以便后续使用
    try:
        with open(TEMP_PATIENT_INFO_FILE, 'w', encoding='utf-8') as f:
            json.dump({
                'patient_name': patient_name,
                'patient_number': patient_number,
                'doctor_name': doctor_name
            }, f, ensure_ascii=False)
    except Exception as e:
        print(f"保存患者信息到文件失败: {e}")
    
    # 重定向到首页进行检测
    return redirect(url_for('home'))

# 清除患者信息的路由
@app.route('/clear_patient_info')
def clear_patient_info():
    if 'patient_info' in session:
        del session['patient_info']
    
    # 清除临时文件
    try:
        if os.path.exists(TEMP_PATIENT_INFO_FILE):
            os.remove(TEMP_PATIENT_INFO_FILE)
    except Exception as e:
        print(f"删除临时患者信息文件失败: {e}")
    
    return redirect(url_for('home'))


# 关于糖尿病页面
@app.route('/about')
def about():
    return render_template('about.html')


# 糖尿病危害页面
@app.route('/dangers')
def dangers():
    return render_template('dangers.html')


# 预防措施页面
@app.route('/prevention')
def prevention():
    return render_template('prevention.html')


# 预测API
@app.route('/predict', methods=['POST'])
def predict_api():
    if 'image' not in request.files:
        return jsonify({'error': '没有文件上传'})

    file = request.files['image']

    if file.filename == '':
        return jsonify({'error': '没有选择文件'})

    if file and allowed_file(file.filename):
        # 保存文件
        upload_path = os.path.join('uploads', file.filename)
        filename = os.path.join(app.config['UPLOAD_FOLDER'], file.filename)
        file.save(filename)

        # 打开并预处理图像
        image = Image.open(filename).convert("RGB")

        # 进行预测
        diagnosis, confidence, all_probs = predict(image)

        # 保存检测记录到数据库
        # 读取图像数据用于存储
        with open(filename, 'rb') as f:
            image_data = f.read()
        
        record_id = save_detection_record(
            filename=file.filename,
            image_path=upload_path,  # 只使用uploads/文件名的格式
            diagnosis=diagnosis,
            confidence=confidence,
            all_probs=all_probs,
            image_data=image_data
        )

        # 准备预测结果
        result = {
            'diagnosis': diagnosis,
            'confidence': float(confidence),
            'all_probs': all_probs,
            'image_path': upload_path,  # 只使用uploads/文件名的格式
            'record_id': record_id
        }

        return jsonify(result)

    return jsonify({'error': '不支持的文件类型'})


# 摄像头预测API
@app.route('/predict_camera', methods=['POST'])
def predict_camera():
    try:
        # 从请求中获取图像数据
        data = request.json
        image_data = data['image'].split(',')[1]  # 移除data:image/png;base64,

        # 将base64数据转换为PIL图像
        image = Image.open(io.BytesIO(base64.b64decode(image_data))).convert("RGB")

        # 进行预测
        diagnosis, confidence, all_probs = predict(image)

        # 保存摄像头拍摄的图像到文件
        timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
        filename = f"camera_{timestamp}.png"
        upload_path = os.path.join('uploads', filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        
        # 保存图像
        image.save(filepath)
        
        # 读取保存的图像数据用于存储
        with open(filepath, 'rb') as f:
            saved_image_data = f.read()
        
        # 保存检测记录到数据库
        record_id = save_detection_record(
            filename=filename,
            image_path=upload_path,  # 只使用uploads/文件名的格式
            diagnosis=diagnosis,
            confidence=confidence,
            all_probs=all_probs,
            image_data=saved_image_data
        )

        # 准备预测结果
        result = {
            'diagnosis': diagnosis,
            'confidence': float(confidence),
            'all_probs': all_probs,
            'image_path': upload_path,  # 只使用uploads/文件名的格式
            'record_id': record_id
        }

        return jsonify(result)
    except Exception as e:
        return jsonify({'error': str(e)})


# 历史记录页面，支持搜索和分页
@app.route('/history')
def history():
    search_term = request.args.get('search', '').strip()
    
    # 获取页码，默认为1
    try:
        page = int(request.args.get('page', 1))
        if page < 1:
            page = 1
    except ValueError:
        page = 1
    
    # 每页显示的记录数
    per_page = 10
    
    # 获取分页记录
    records, total_count = get_all_records(search_term, page, per_page)
    
    # 计算总页数
    total_pages = (total_count + per_page - 1) // per_page
    
    # 确保当前页码不超过总页数
    if page > total_pages and total_pages > 0:
        page = total_pages
        # 重新获取正确页码的记录
        records, total_count = get_all_records(search_term, page, per_page)
    
    return render_template('history.html', 
                          records=records, 
                          search_term=search_term, 
                          total_count=total_count,
                          page=page,
                          total_pages=total_pages,
                          per_page=per_page)

# 删除记录
@app.route('/history/delete/<int:record_id>')
def delete_record_route(record_id):
    success = delete_record(record_id)
    if success:
        flash('记录已成功删除', 'success')
    else:
        flash('删除失败，记录不存在', 'danger')
    return redirect(url_for('history'))

# 历史记录详情页面
@app.route('/history/<int:record_id>')
def record_detail(record_id):
    record = get_record_detail(record_id)
    if record:
        return render_template('record_detail.html', record=record)
    else:
        flash('记录不存在')
        return redirect(url_for('history'))

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0')